#!/usr/bin/python
#
# (c) 2006 Travis F Vachon
#
# This software is distributable under the terms of the GNU
# General Public License (GPL) v2, the text of which can be found at
# http://www.gnu.org/copyleft/gpl.html. Installing, importing or otherwise
# using this module constitutes acceptance of the terms of this License.
#
# Disclaimer
# 
# This software is provided "as-is".  There are no expressed or implied
# warranties of any kind, including, but not limited to, the warranties
# of merchantability and fittness for a given application.  In no event
# shall Travis Vachon be liable for any direct, indirect, incidental,
# special, exemplary or consequential damages (including, but not limited
# to, loss of use, data or profits, or business interruption) however
# caused and on any theory of liability, whether in contract, strict
# liability or tort (including negligence or otherwise) arising in any way
# out of the use of this software, even if advised of the possibility of
# such damage.
#

import sys,os
import getopt
import xen.lowlevel.xc as xc
import logging
import threading
import time
from conf import *
import globals
import socket
from xen.xend.XendClient import server as xen_server

if __name__ == "__main__":

    #parse options
    
    opts, args = getopt.getopt(sys.argv[1:],"l:")
    opts = dict(opts)
    logfilename = opts.get("-l","/var/log/xenbal.log")

logfilename = "/var/log/xenbal.log"


#Set up logging facilities

logger = logging.getLogger('xenbald')
hdlr = logging.FileHandler(logfilename)
formatter = logging.Formatter('[%(asctime)s %(name)s] %(levelname)s (%(module)s:%(lineno)d) %(message)s')
hdlr.setFormatter(formatter)
logger.addHandler(hdlr)
logger.setLevel(logging.DEBUG)

class MigFuncNotDefinedException(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)

class MigrationFunction(object):
    def __init__(self, args=None):
        if args == None: args == None
        self.args = args

    def get_migration_list(self):
        raise MigFuncNotDefinedException, "Migration Function Not Defined. Please define a migration list function."

    def get_number_bal_dest(self):
        num_domains = len(globals.this_node.domains)

        for node in globals.peer_nodes:
            if len(node.domains) < (num_domains - 1):
                return node

        return None

class NobalFunction(MigrationFunction):
    def get_migration_list(self):
        return [{'domain':None,'destination':None}]


class SimpleSIFunction(MigrationFunction):
    def __init__(self, args): 
        self.pause = args.get('pause',60)
        self.live = args.get('live',1)
        self.load_memory = args.get('load_memory',1)
        self.load_threshold = args.get('load_threshold',0.95)

        MigrationFunction.__init__(self, args)

    def get_number_bal_dest(self):
        num_domains = len(globals.this_node.domains)

        for node in globals.peer_nodes.values():
            if len(node.domains) < (num_domains - 1):
                return node

        return None


    def get_migration_list(self):
        logger.info("SimpleSI Function")
        if globals.this_node.get_cpu_load(self.load_memory) < self.load_threshold:
            logger.debug("Underloaded, not migrating")
            return [{'domain':None,'destination':None}]
        elif len(globals.this_node.domains) < 2:
            logger.debug("Only one domain, not migrating")
            return [{'domain':None,'destination':None}]
        elif globals.this_node.get_min_load_peer(self.load_memory)[1] < 0.5:
            logger.info(str(globals.this_node.get_min_load_peer(self.load_memory)[1]))
            return [{'domain':globals.this_node.get_max_load_domain(self.load_memory)[0],
                     'destination':globals.this_node.get_min_load_peer(self.load_memory)[0],
                     'live':self.live,
                     'pause':self.pause}]
        else:
            number_bal_dest = self.get_number_bal_dest()
            logger.info(str(number_bal_dest))
            if not number_bal_dest == None:
                return [{'domain':globals.this_node.get_max_load_domain(self.load_memory)[0],
                         'destination':number_bal_dest,
                         'live':self.live,
                         'pause':self.pause}]

            
        return [{'domain':None,'destination':None}]

class SimpleRIFunction(MigrationFunction):
    def __init__(self, args):
        self.pause = args.get('pause',60)
        self.live = args.get('live',1)
        self.load_memory = args.get('load_memory',1)
        self.load_threshold = args.get('load_threshold',0.3)
        MigrationFunction.__init__(self, args)

    def get_number_bal_source(self):
        num_domains = len(globals.this_node.domains)

        for node in globals.peer_nodes.values():
            if len(node.domains) > (num_domains + 1):
                return node

        return None


    def get_migration_list(self):

        if globals.this_node.get_cpu_load(self.load_memory) > self.load_threshold:
            number_bal_source = self.get_number_bal_source()
            if not number_bal_source == None:
                
                return [{'domain':number_bal_source.get_max_load_domain(self.load_memory)[0],
                         'destination':globals.this_node,
                         'source':number_bal_source,
                         'live':self.live,
                         'pause':self.pause}]

            logger.debug("Overloaded, not migrating")
            return [{'domain':None,'destination':None}]
        
        else:
            chosen_source = None
            peer_list = globals.this_node.get_sorted_peer_load_list()
            peer_list.reverse()
            for peer in peer_list:
                if len(peer.domains) > 1 and peer.get_cpu_load(self.load_memory) > 0.8:
                    chosen_source = peer
                    break
            else:
                # We couldn't find a good source
                number_bal_source = self.get_number_bal_source()
                if not number_bal_source == None:
                    
                    return [{'domain':number_bal_source.get_max_load_domain(self.load_memory)[0],
                             'destination':globals.this_node,
                             'source':number_bal_source,
                             'live':self.live,
                             'pause':self.pause}]

                else:
                    return [{'domain':None,'destination':None}]
                
            return [{'domain':chosen_source.get_max_load_domain(self.load_memory)[0],
                         'destination':globals.this_node,
                         'source':chosen_source,
                         'live':self.live,
                         'pause':self.pause}]


            
        
  
    
class UtilityFunction(MigrationFunction):
    '''Define loss as the lost potential for work to be done.
    Consider a fully loaded 1Ghz node with two domains each receiving
    50% of the cputime. We will assume in
    this case that each of the nodes would be capable of using the
    entire processor (since we cannot tell instantaneously (that is,
    without more knowledge of the system than we can glean from a snapshot)
    how much it would ideally have. The loss of a domain is the amount
    of the processor it is capable of using minus the amount of the
    processor it is currently using. Then the loss of each domain
    in this case is 1Ghz - 0.5Ghz = 0.5Ghz. The total loss of the system
    is 1 Ghz.
    
    Define unfairness of the system as the difference in loss between the
    domain with the largest loss and the domain with the smallest loss.
    The in the system above, the unfairness is equal to 0.
    Consider two fully loaded 1Ghz nodes with 3 domains on the first and
    1 domain on the second. The loss of the 3 domains on the first node is
    1-0.33 = 0.66Ghz each, while the loss of the 1 domain on the second
    node is 1-1 = 0Ghz. The unfairness of the system is 0.66Ghz.
    
    This algorithm will minimize unfairness and total loss. 
    '''
    
    def __init__(self, args):

        MigrationFunction.__init__(self, args)
        
        
        self.load_memory = args.get('load_memory',10)
        self.sensitivity = args.get('sensitivity',0.8)
        self.pause = args.get('pause',60)
        self.live = args.get('live',1)


        self.current_total_loss = None
        self.current_unfairness = None

        self.best_loss_so_far = None
        self.best_unfairness_so_far = None

        self.best_loss_scenario = None
        self.best_unfairness_scenario = None


    def update_loss_statistics(self,dom, dest, loss, unfairness, source=globals.this_node):
        
        if self.best_loss_so_far == None or loss < self.best_loss_so_far:
            self.best_loss_so_far = loss
            self.best_loss_scenario['domain'] = dom
            self.best_loss_scenario['destination'] = dest
            self.best_loss_scenario['source'] = source
        if self.best_unfairness_so_far == None or unfairness < self.best_unfairness_so_far:
            self.best_unfairness_so_far = unfairness
            self.best_unfairness_scenario['domain'] = dom
            self.best_unfairness_scenario['destination'] = dest
            self.best_unfairness_scenario['source'] = source


    def reset_statistics(self):
        #Calculate the current loss in the system
        current_total_loss_stats = globals.this_node.simulate_migration(None,None, self.load_memory)
        self.current_total_loss = current_total_loss_stats['loss']
        self.current_unfairness = current_total_loss_stats['unfairness']


        self.best_loss_so_far = None
        self.best_unfairness_so_far = None

        self.best_loss_scenario = {'domain':None, 'destination':None, 'live':self.live}
        self.best_unfairness_scenario = {'domain':None, 'destination':None, 'live':self.live}

            

                
    def pick_migration(self):
        if ((not self.best_loss_scenario['domain'] == None) and
            (not self.best_loss_scenario['destination'] == None)):
            
            if self.best_loss_so_far < self.sensitivity*self.current_total_loss:
                self.best_loss_scenario['pause'] = self.pause
                print 
                print "Loss migration"
                print
                return [self.best_loss_scenario]
            
            if ((not self.best_unfairness_scenario['domain'] == None) and
                (not self.best_unfairness_scenario['destination'] == None)):
                    
                if self.best_unfairness_so_far < self.sensitivity*self.current_unfairness:
                    print
                    print "UF migration"
                    print
                    self.best_unfairness_scenario['pause'] = self.pause
                    return [self.best_unfairness_scenario]

            #If we get here, we want to do nothing
            return [{'domain':None,'destination':None}]


class UtilityFunctionSI(UtilityFunction):
    def __init__(self,args=None):
        if args == None: args == None
        UtilityFunction.__init__(self,args)

    def get_migration_list(self):
        migration_dictionary_list = []
        
        self.reset_statistics()
        
        for node in globals.peer_nodes.values():
            for domain in globals.this_node.domains.values():
                if not str(domain.id) == '0':
                    
                    loss_stats = globals.this_node.simulate_migration(domain, node, self.load_memory)
                    self.update_loss_statistics(domain, node,loss_stats['loss'],loss_stats['unfairness'])
                    
        return self.pick_migration()

class UtilityFunctionRI(UtilityFunction):
    def __init__(self,args=None):
        if args == None: args == None
        UtilityFunction.__init__(self,args)

    def get_migration_list(self):
        migration_dictionary_list = []
        
        self.reset_statistics()
        
        for node in globals.peer_nodes.values():
            for domain in node.domains.values():
                if not str(domain.id) == '0':
                    
                    loss_stats = node.simulate_migration(domain, globals.this_node, self.load_memory)
                    self.update_loss_statistics(domain, globals.this_node, loss_stats['loss'],loss_stats['unfairness'], source=node)
                    
        return self.pick_migration()


# The Balancer.
NO_BAL_POLICY = 0
SENDER_INIT_POLICY = 1
RECEIVER_INIT_POLICY = 2
SYM_INIT_POLICY = 3
class LoadBalancer(threading.Thread):
    def __init__(self,policy=NO_BAL_POLICY, policy_args = None):
        threading.Thread.__init__(self)

        self.peer_list = globals.peer_list

        print "Querying " + str(self.peer_list)
        self.set_no_bal_policy()



    
    def set_simplesi_policy(self, args):
        logger.info("Using SimpleSIFunction")
        self.policy_function = SimpleSIFunction(args)
        return True

    def set_simpleri_policy(self, args):
        self.policy_function = SimpleRIFunction(args)
        return True

    def set_utilitysi_policy(self, args):
        self.policy_function = UtilityFunctionSI(args)
        return True

    def set_utilityri_policy(self, args):
        self.policy_function = UtilityFunctionRI(args)
        return True

    def set_no_bal_policy(self):
        self.policy_function = NobalFunction()
        return True
    

    def run(self):
        self.run = True
        while(self.run):

            try:
                migration_dict_list = self.policy_function.get_migration_list()
                logger.debug(str(migration_dict_list))
                #Migrate here...
                if not migration_dict_list == None:
                    for migration_dict in migration_dict_list:
                        chosen_source = migration_dict.get('source',globals.this_node)
                        if (not migration_dict['domain'] == None and
                            not migration_dict['destination'] == None and
                            not chosen_source == None):

                            chosen_domain = migration_dict['domain'].id
                            chosen_destination = migration_dict['destination'].name

                            live = migration_dict.get('live',1)
                            resources = migration_dict.get('resources',0)
                            pause = migration_dict.get('pause',30)

                            logger.info("migrating " + str(chosen_domain) + " to " + chosen_destination)
                            try:
                                logger.debug("migrating")
                                chosen_source.server.migrate(chosen_domain, chosen_destination, 1,0)
                                logger.debug("migrated")                            
                                time.sleep(pause)
                            except:
                                logger.warning("Migrate failed!")
                        else:
                            logger.info("Not migrating: dom="+str(migration_dict['domain'])+" dest="+str(migration_dict['destination']))
                        
            finally:
                    time.sleep(1)



    def exit(self):
        self.run = False;

